# 清华文本分类数据集
import os
import random
import hashlib
import pandas as pd
from sklearn.naive_bayes import GaussianNB
from sklearn.model_selection import train_test_split


def txt2sha512(txt):
    m = hashlib.sha512()
    m.update(txt.encode('utf-8'))
    sha512 = m.hexdigest()
    return sha512


class thuc:
    def __init__(self):
        self.path = "D:\Files\Downloads\THUCNews"

        self.labels = []
        self.label_path = {}
        self.max_t = 0

        self.df_txt = None
        self.df_label = None
        self.df_test = None
        self.df_test_label = None

        self.count = 0

    # 读取目录并生成标签路径对
    def read_dir(self):

        self.labels = os.listdir(self.path)
        for label in self.labels:
            l_path = os.path.join(self.path, label)
            t_s = os.listdir(l_path)
            for ti in t_s:
                if label not in self.label_path.keys():
                    self.label_path[label] = [os.path.join(l_path, ti)]
                else:
                    self.label_path[label].append(os.path.join(l_path, ti))

                c_len = len(self.label_path[label])
                if c_len > self.max_t:
                    self.max_t = c_len

    # 打乱顺序
    def rand_shuffle(self):
        for label in self.label_path.keys():
            random.shuffle(self.label_path[label])

    # 读取文本内容
    def read_text(self):
        self.df_txt = pd.DataFrame(columns=[str(i) for i in range(128)])
        self.df_label = pd.DataFrame(columns=['label'])

        for i in range(self.max_t):
            for l in self.labels:
                if i < len(self.label_path[l]):
                    c_path = self.label_path[l][i]

                    with open(c_path, 'r', encoding='utf-8') as fr:
                        c_hash = txt2sha512(fr.read())
                        c_l = self.labels.index(l)

                        self.c_data(c_hash, c_l)
                        self.count += 1
                        if self.count % 100 == 0:
                            print(self.count)

                        if self.count % 10000 == 0:
                            return

    def get_data(self):

        self.df_txt, self.df_test, self.df_label, self.df_test_label = train_test_split(self.df_txt, self.df_label, test_size=0.3, random_state=0)

        self.df_txt = self.df_txt.astype(str)
        self.df_test = self.df_test.astype(str)
        self.df_label = self.df_label.astype(int)
        self.df_test_label = self.df_test_label.astype(int)
        print(self.df_txt)
        print(self.df_label)
        print(self.df_test)
        print(self.df_test_label)

    def c_data(self, c_hash, c_l):
        c_d = {}
        for i in range(128):
            c_d[str(i)] = [int(c_hash[i], 16)]

        self.df_txt = pd.concat([self.df_txt, pd.DataFrame(c_d)], ignore_index=True)
        self.df_label = pd.concat([self.df_label, pd.DataFrame({'label': [c_l]})], ignore_index=True)


class bayes_model:
    def __int__(self):
        pass

    def train_model(self, train_x, train_y):
        clf = GaussianNB()
        clf.fit(train_x, train_y)
        return clf

    def proba_data(self, clf, test_x, test_y):
        y_predict = clf.predict(test_x)
        count = 0
        true = 0
        false = 0
        test_y = test_y['label'].values.tolist()
        for i in range(len(y_predict)):
            count += 1
            if test_y[i] == y_predict[i]:
                true += 1
            else:
                false += 1

        print(count)
        print(true)
        print(false)
        return

    def exc_p(self, train_x, test_x, train_y, test_y):
        clf = self.train_model(train_x, train_y)
        self.proba_data(clf, test_x, test_y)
        return


if __name__ == '__main__':
    t = thuc()
    t.read_dir()
    t.rand_shuffle()
    t.read_text()
    t.get_data()

    b = bayes_model()
    b.exc_p(t.df_txt, t.df_test, t.df_label, t.df_test_label)
